# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

import logging
import time
import threading
from concurrent.futures import ThreadPoolExecutor

import tornado.gen
import requests
from google.protobuf import json_format
from absl import app

import coin.exchange.base.kr_rest.private_client_base as privcb
from coin.exchange.uniswap_v3.kr_rest.product import UniswapProduct
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from xunkemgmt_client.client.api_client import XunkemgmtClient
from coin.exchange.uniswap_v2.kr_rest.private_client_v2 import UniswapPrivateParser
from coin2.service.web3.order_service_pb2 import (
    QueryAccountBalanceRequestProto, QueryAccountBalanceResponseProto,
    QueryNftTokenInfoRequestProto, QueryNftTokenInfoResponseProto,
    QueryProvideLiquidityRequestProto, QueryProvideLiquidityResponseProto,
    QuerySubmitOrderRequestProto, QuerySubmitOrderResponseProto,
    QueryWithdrawLiquidityRequestProto, QueryWithdrawLiquidityResponseProto, BalanceItem
)
from coin.exchange.uniswap_v3.kr_rest.constants import currency_list
from coin.blockchain.request_util import parse_coin_key


class UniswapPrivateClient(privcb.PrivateClientBase):
    ProductType = UniswapProduct

    def __init__(
        self,
        wallet_address,
        og_info,
        blockchain_name="ethereum",
        mea="Spot.Uniswap.v3",
    ):
        privcb.PrivateClientBase.__init__(self)
        self._wallet_address = wallet_address
        self._blockchain_name = blockchain_name
        self._mea = mea
        self._currency_list = [curr.native_currency for curr in og_info.currencies]
        self._symbol_list = [prod.native_symbol for prod in og_info.products]
        self._web3_api_client = XunkemgmtClient()
        self._executor = ThreadPoolExecutor(max_workers=2)
        self._fire_lock_map = {symbol: threading.Lock() for symbol in self._symbol_list}

    def query_account_balance_impl(self):
        raise NotImplementedError()

    def _async_account_balance_impl(self):
        if self._executor:
            balance_request = QueryAccountBalanceRequestProto(
                blockchain_name=self._blockchain_name,
                mea=self._mea,
                wallet_address=self._wallet_address,
                currencies=self._currency_list,
            )
            fut = self._executor.submit(self._web3_api_client.query_account_balance, balance_request)
            return fut
        return None

    @tornado.gen.coroutine
    def async_query_account_balance_impl(self):
        req_timestamp = time.time()
        fut = self._async_account_balance_impl()
        account_balance_response_proto = yield fut
        resp_timestamp = time.time()
        update = RestUpdate(
            req_timestamp,
            resp_timestamp,
            account_balance_response_proto,
            requests.codes.ok,
        )
        update.msg = UniswapPrivateParser.parse_balance(update.msg)
        return update

    def _async_query_nft_token_info(self, symbol: str):
        if self._executor:
            nft_token_info_request = QueryNftTokenInfoRequestProto(
                blockchain_name=self._blockchain_name,
                mea=self._mea,
                wallet_address=self._wallet_address,
                symbol=symbol,
            )
            fut = self._executor.submit(self._web3_api_client.query_nft_token_info, nft_token_info_request)
            return fut
        return None

    @tornado.gen.coroutine
    def async_query_nft_token_info(self):
        nft_token_infos = {}
        for symbol in self._symbol_list:
            fut = self._async_query_nft_token_info(symbol)
            nft_token_info_response_proto = yield fut
            nft_token_infos[symbol] = json_format.MessageToDict(
                nft_token_info_response_proto, preserving_proto_field_name=True)
        return nft_token_infos

    def query_nft_token_info(self):
        nft_token_infos = {}
        for symbol in self._symbol_list:
            nft_token_info_request = QueryNftTokenInfoRequestProto(
                blockchain_name=self._blockchain_name,
                mea=self._mea,
                wallet_address=self._wallet_address,
                symbol=symbol,
            )
            nft_token_info_response_proto = self._web3_api_client.query_nft_token_info(nft_token_info_request)
            nft_token_infos[symbol] = json_format.MessageToDict(
                nft_token_info_response_proto, preserving_proto_field_name=True)
        return nft_token_infos

    def _withdraw_liquidity(self, symbol, execution_lock):
        withdraw_liquidity_request = QueryWithdrawLiquidityRequestProto(
            blockchain_name=self._blockchain_name,
            mea=self._mea,
            wallet_address=self._wallet_address,
            symbol=symbol,
        )
        try:
            withdraw_liquidity_response_proto = self._web3_api_client.query_withdraw_liquidity(
                withdraw_liquidity_request)
        except Exception:
            logging.exception('Exception raised from _web3_api_client.query_withdraw_liquidity.')
            if execution_lock and execution_lock.locked():
                execution_lock.release()
            return None
        if execution_lock:
            assert execution_lock.locked()
            execution_lock.release()
        return json_format.MessageToDict(withdraw_liquidity_response_proto, preserving_proto_field_name=True)

    def _provide_liquidity(
        self,
        symbol,
        execution_lock,
        price_upper,
        price_lower,
        **kwargs
    ):
        provide_liquidity_request = QueryProvideLiquidityRequestProto(
            blockchain_name=self._blockchain_name,
            mea=self._mea,
            wallet_address=self._wallet_address,
            symbol=symbol,
            min_price=price_lower,
            max_price=price_upper,
        )
        try:
            provide_liquidity_response_proto = self._web3_api_client.query_provide_liquidity(
                provide_liquidity_request)
        except Exception:
            logging.exception('Exception raised from _web3_api_client.query_provide_liquidity.')
            if execution_lock and execution_lock.locked():
                execution_lock.release()
            return None
        if execution_lock:
            assert execution_lock.locked()
            execution_lock.release()
        return json_format.MessageToDict(provide_liquidity_response_proto, preserving_proto_field_name=True)

    def query_liquidity_pool_submit_impl(self, product, action_type, execution_lock, **kwargs):
        if self._executor:
            if action_type == "PROVIDE":
                return self._executor.submit(self._provide_liquidity, product.native_symbol, execution_lock, **kwargs)
            elif action_type == "WITHDRAW":
                return self._executor.submit(self._withdraw_liquidity, product.native_symbol, execution_lock)
            else:
                raise ValueError(product, action_type)
        return None

    def _query_submit_impl(self, symbol, side, price, qty):
        lock = self._fire_lock_map[symbol]
        if lock.acquire(False):
            order_request = QuerySubmitOrderRequestProto(
                blockchain_name=self._blockchain_name,
                mea=self._mea,
                wallet_address=self._wallet_address,
                symbol=symbol,
                price=float(price),
                qty=float(qty),
                side=side.name,
            )
            try:
                order_response_proto = self._web3_api_client.query_submit_order(order_request)
            except Exception:
                logging.exception('Exception raised from _web3_api_client.query_submit_order.')
                lock.release()
                return None
            lock.release()
            return json_format.MessageToDict(order_response_proto, preserving_proto_field_name=True)
        else:
            logging.warning(f"order rejected: {symbol}")
            return None

    def query_submit_impl(self, product, side, price, qty):
        symbol = product.native_symbol
        if self._executor:
            return self._executor.submit(self._query_submit_impl, symbol, side, price, qty)
        else:
            return None

    def query_cancel_product_impl(self, product):
        raise NotImplementedError()

    def query_fills_impl(self, product):
        raise NotImplementedError()

    def query_list_orders_impl(self, product):
        raise NotImplementedError()


class UniswapWalletPrivateClient(privcb.PrivateClientBase):
    def __init__(
        self,
        key_file,
        mea="Spot.Uniswap.v3-arbitrum",
    ):
        privcb.PrivateClientBase.__init__(self)
        self._key = parse_coin_key(key_file).chain_keys[0]
        self._wallet_address = self._key.access_key
        self._blockchain_name = self._key.blockchain
        self._mea = mea
        self._currency_list = currency_list
        self._web3_api_client = XunkemgmtClient()

    def _combine_eth(self, account_balance_response):
        balance_dict = dict(account_balance_response.currency_to_balance)
        weth_balance = balance_dict['WETH']
        eth_balance = balance_dict['ETH']
        balance_dict['ETH'] = BalanceItem(
            available = weth_balance.available + eth_balance.available,
            hold = weth_balance.hold + eth_balance.hold,
        )
        del balance_dict['WETH']
        resp = QueryAccountBalanceResponseProto(
            status = account_balance_response.status,
            wallet_address = account_balance_response.wallet_address,
            currency_to_balance = balance_dict
        )
        return resp

    def query_account_balance_impl(self):
        req_timestamp = time.time()
        balance_request = QueryAccountBalanceRequestProto(
            blockchain_name=self._blockchain_name,
            mea=self._mea,
            wallet_address=self._wallet_address,
            currencies=self._currency_list,
        )
        account_balance_response_proto = self._web3_api_client.query_account_balance(balance_request)
        account_balance_response_proto = self._combine_eth(account_balance_response_proto)
        resp_timestamp = time.time()
        update = RestUpdate(
            req_timestamp,
            resp_timestamp,
            account_balance_response_proto,
            requests.codes.ok,
        )
        update.msg = UniswapPrivateParser.parse_balance(update.msg)
        return update

    def query_cancel_product_impl(self, product):
        raise NotImplementedError()

    def query_fills_impl(self, product):
        raise NotImplementedError()

    def query_list_orders_impl(self, product):
        raise NotImplementedError()

    def query_submit_impl(self, product, side, price, qty):
        raise NotImplementedError()


def main(_):
    wallet_file = '../../coin_key/view/xunke11/uniswap_arbitrum/view_key.json'
    private_client = UniswapWalletPrivateClient(
        wallet_file, "Spot.Uniswap.v3-arbitrum")
    print(private_client.query_account_balance_impl())


if __name__ == '__main__':
    app.run(main)
